{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import hail as hl\n",
    "hl.init(spark_conf={\"spark.hadoop.fs.gs.requester.pays.mode\": \"AUTO\",\n",
    "                    \"spark.hadoop.fs.gs.requester.pays.project.id\": \"broad-ctsa\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The workflow I ended up using to add these GIANT didn't fit nicely into the old extract/load setup, so I added this notebook directory.\n",
    "\n",
    "For this particular case, the files were small enough that I was exploring things locally, and it ended up being easiest to just get things done here. You can run the cells for height/bmi/whr independently to generate the tables and entries for `datasets.json`.\n",
    "\n",
    "It just assumes that you have the [GIANT consortium data files](https://portals.broadinstitute.org/collaboration/giant/index.php/GIANT_consortium_data_files) downloaded to a `tmp` folder in your home directory , e.g. for height, `~/tmp/giant_height_exome_summary/` would contain all the `.txt` files for height."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Open our datasets config file so we can add our newly generated entries\n",
    "datasets_path = os.path.abspath(\"../../hail/python/hail/experimental/datasets.json\")\n",
    "with open(datasets_path, \"r\") as f:\n",
    "    datasets = json.load(f)\n",
    "\n",
    "# To generate correct descriptions in datasets.json entries for each population\n",
    "json_populations = dict(\n",
    "    zip(\n",
    "        [\"AFR\", \"EUR\", \"EAS\", \"AMR\", \"SAS\", \"ALL\"],\n",
    "        [\"African/African-American\", \"European\", \"East Asian\", \"Latino/Admixed American\", \"South Asian\", \"all\"]\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Height"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Iterate through GIANT height text files to write Hail Tables to GCS and generate entries to insert into datasets.json:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# For GIANT 2018 Exome Array Summary Statistics - Height\n",
    "populations = [\"AA\", \"EA\", \"EAS\", \"HA\", \"SA\", \"All\"]\n",
    "renamed_populations = [\"AFR\", \"EUR\", \"EAS\", \"AMR\", \"SAS\", \"ALL\"]\n",
    "\n",
    "# Map population name to maf and exac_maf field names\n",
    "maf = dict(\n",
    "    zip(populations, [\"AFR_MAF\", \"EUR_MAF\", \"EAS_MAF\", \"AMR_MAF\", \"SAS_MAF\", \"GMAF\"])\n",
    ")\n",
    "exac_maf = dict(\n",
    "    zip(populations, [\"ExAC_AFR_MAF\", \"ExAC_NFE_MAF\", \"ExAC_EAS_MAF\", \"ExAC_AMR_MAF\", \"ExAC_SAS_MAF\", \"ExAC_MAF\"])\n",
    ")\n",
    "\n",
    "# Re-map population names for file name of Hail Table\n",
    "output_name = dict(zip(populations, renamed_populations))\n",
    "\n",
    "for population in populations:\n",
    "    print(population)\n",
    "    \n",
    "    input_file = os.path.expanduser(\"~\") + f\"/tmp/giant_height_exome_summary/height_{population}_add_SV.txt\"\n",
    "    name = f\"giant_height_exome_{output_name[population]}\"\n",
    "    version = \"2018\"\n",
    "    build = \"GRCh37\"\n",
    "    \n",
    "    ht = hl.import_table(input_file,\n",
    "                         impute=True,\n",
    "                         missing=[\"-\", \"NA\", \"Inf\"],\n",
    "                         delimiter=\"\\s+\",\n",
    "                         types = {\"beta\": hl.tfloat64,\n",
    "                                  \"se\": hl.tfloat64,\n",
    "                                  \"Pvalue\": hl.tfloat64})\n",
    "\n",
    "    ht2 = ht.annotate(locus = hl.locus(ht.CHR, ht.POS, reference_genome=build),\n",
    "                      alleles = [ht.REF, ht.ALT],\n",
    "                      temp_maf = hl.dict(\n",
    "                          ht[maf[population]].split(\",\").map(\n",
    "                              lambda x: (x.split(\":\")[0], hl.float(x.split(\":\")[1]))\n",
    "                          )\n",
    "                      ),\n",
    "                      temp_exac_maf = hl.dict(\n",
    "                          ht[exac_maf[population]].split(\",\").map(\n",
    "                              lambda x: (x.split(\":\")[0], hl.float(x.split(\":\")[1]))\n",
    "                          )\n",
    "                      )\n",
    "                     )\n",
    "    \n",
    "    ht2 = ht2.select(\"locus\", \"alleles\", \"SNPNAME\", \"temp_maf\", \"temp_exac_maf\", \"beta\", \"se\", \"Pvalue\")\n",
    "    ht2 = ht2.rename({\"temp_maf\" : maf[population].lower(),\n",
    "                      \"temp_exac_maf\" : exac_maf[population].lower(),\n",
    "                      \"SNPNAME\": \"snp_name\",\n",
    "                      \"Pvalue\": \"pvalue\"})\n",
    "    ht2 = ht2.key_by(\"locus\", \"alleles\")\n",
    "\n",
    "    n_rows = ht2.count()\n",
    "    n_partitions = ht2.n_partitions()\n",
    "\n",
    "    ht2 = ht2.annotate_globals(metadata=hl.struct(name=name,\n",
    "                                                  version=version,\n",
    "                                                  reference_genome=build,\n",
    "                                                  n_rows=n_rows,\n",
    "                                                  n_partitions=n_partitions))\n",
    "    \n",
    "    for region in [\"us\", \"eu\"]:\n",
    "        output_file = f\"gs://hail-datasets-{region}/{name}_{version}_{build}.ht\"\n",
    "        ht2.write(output_file, overwrite=True)\n",
    "\n",
    "    json_entry = {\n",
    "        \"annotation_db\": {\n",
    "            \"key_properties\": [\n",
    "                \"unique\"\n",
    "            ]\n",
    "        },\n",
    "        \"description\": f\"GIANT (Genetic Investigation of ANthropometric Traits): \"\n",
    "                       f\"height exome array summary statistics Hail Table for {json_populations[output_name[population]]} \"\n",
    "                       f\"population(s).\",\n",
    "        \"url\": \"https://portals.broadinstitute.org/collaboration/giant/index.php/GIANT_consortium_data_files\",\n",
    "        \"versions\": [\n",
    "            {\n",
    "                \"reference_genome\": build,\n",
    "                \"url\": {\n",
    "                    \"aws\": {\n",
    "                        \"us\": f\"s3://hail-datasets-us-east-1/{name}_{version}_{build}.ht\"\n",
    "                    },\n",
    "                    \"gcp\": {\n",
    "                        \"us\": f\"gs://hail-datasets-us/{name}_{version}_{build}.ht\",\n",
    "                        \"eu\": f\"gs://hail-datasets-eu/{name}_{version}_{build}.ht\"\n",
    "                    }\n",
    "                },\n",
    "                \"version\": version\n",
    "            }\n",
    "        ]\n",
    "    }\n",
    "    datasets[name] = json_entry\n",
    "\n",
    "# Write new entries back to datasets.json config:\n",
    "with open(datasets_path, \"w\") as f:\n",
    "    json.dump(datasets, f, sort_keys=True, ensure_ascii=False, indent=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### BMI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Iterate through GIANT body mass index (BMI) text files to write Hail Tables to GCS and generate entries to insert into datasets.json:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# For GIANT 2018 Exome Array Summary Statistics - BMI\n",
    "populations = [\"African_American\", \"European\", \"Eastern_Asian\", \"Hispanic_American\", \"South_Asian\", \"All_ancestry\"]\n",
    "renamed_populations = [\"AFR\", \"EUR\", \"EAS\", \"AMR\", \"SAS\", \"ALL\"]\n",
    "\n",
    "# Map population name to maf and exac_maf field names\n",
    "maf = dict(\n",
    "    zip(populations, [\"AFR_MAF\", \"EUR_MAF\", \"EAS_MAF\", \"AMR_MAF\", \"SAS_MAF\", \"GMAF\"])\n",
    ")\n",
    "exac_maf = dict(\n",
    "    zip(populations, [\"ExAC_AFR_MAF\", \"ExAC_NFE_MAF\", \"ExAC_EAS_MAF\", \"ExAC_AMR_MAF\", \"ExAC_SAS_MAF\", \"ExAC_MAF\"])\n",
    ")\n",
    "\n",
    "# Re-map population names for file name of Hail Table\n",
    "output_name = dict(zip(populations, renamed_populations))\n",
    "\n",
    "for population in populations:\n",
    "    print(population)\n",
    "    \n",
    "    input_file = os.path.expanduser(\"~\") + f\"/tmp/giant_bmi_exome_summary/BMI_{population}.fmt.gzip.txt\"\n",
    "    name = f\"giant_bmi_exome_{output_name[population]}\"\n",
    "    version = \"2018\"\n",
    "    build = \"GRCh37\"\n",
    "    \n",
    "    ht = hl.import_table(input_file,\n",
    "                         impute=True,\n",
    "                         missing=[\"-\", \"NA\", \"Inf\"],\n",
    "                         delimiter=\"\\s+\",\n",
    "                         types = {\"beta\": hl.tfloat64,\n",
    "                                  \"se\": hl.tfloat64,\n",
    "                                  \"Pvalue\": hl.tfloat64})\n",
    "\n",
    "    ht2 = ht.annotate(locus = hl.locus(ht.CHR, ht.POS, reference_genome=build),\n",
    "                      alleles = [ht.REF, ht.ALT],\n",
    "                      snp_name = ht.SNPNAME,\n",
    "                      pvalue = ht.Pvalue,\n",
    "                      temp_maf = hl.dict(\n",
    "                          ht[maf[population]].split(\",\").map(\n",
    "                              lambda x: (x.split(\":\")[0], hl.float(x.split(\":\")[1]))\n",
    "                          )\n",
    "                      ),\n",
    "                      temp_exac_maf = hl.dict(\n",
    "                          ht[exac_maf[population]].split(\",\").map(\n",
    "                              lambda x: (x.split(\":\")[0], hl.float(x.split(\":\")[1]))\n",
    "                          )\n",
    "                      )\n",
    "                     )\n",
    "\n",
    "    ht2 = ht2.select(\"locus\", \"alleles\", \"SNPNAME\", \"temp_maf\", \"temp_exac_maf\", \"beta\", \"se\", \"Pvalue\")\n",
    "    ht2 = ht2.rename({\"temp_maf\" : maf[population].lower(),\n",
    "                      \"temp_exac_maf\" : exac_maf[population].lower(),\n",
    "                      \"SNPNAME\": \"snp_name\",\n",
    "                      \"Pvalue\": \"pvalue\"})\n",
    "    ht2 = ht2.key_by(\"locus\", \"alleles\")\n",
    "\n",
    "    n_rows = ht2.count()\n",
    "    n_partitions = ht2.n_partitions()\n",
    "\n",
    "    ht2 = ht2.annotate_globals(metadata=hl.struct(name=name,\n",
    "                                                  version=version,\n",
    "                                                  reference_genome=build,\n",
    "                                                  n_rows=n_rows,\n",
    "                                                  n_partitions=n_partitions))\n",
    "\n",
    "    for region in [\"us\", \"eu\"]:\n",
    "        output_file = f\"gs://hail-datasets-{region}/{name}_{version}_{build}.ht\"\n",
    "        ht2.write(output_file, overwrite=True)\n",
    "\n",
    "    json_entry = {\n",
    "        \"annotation_db\": {\n",
    "            \"key_properties\": [\n",
    "                \"unique\"\n",
    "            ]\n",
    "        },\n",
    "        \"description\": f\"GIANT (Genetic Investigation of ANthropometric Traits): \"\n",
    "                       f\"body mass index (BMI) exome array summary statistics Hail Table \"\n",
    "                       f\"for {json_populations[output_name[population]]} population(s).\",\n",
    "        \"url\": \"https://portals.broadinstitute.org/collaboration/giant/index.php/GIANT_consortium_data_files\",\n",
    "        \"versions\": [\n",
    "            {\n",
    "                \"reference_genome\": build,\n",
    "                \"url\": {\n",
    "                    \"aws\": {\n",
    "                        \"us\": f\"s3://hail-datasets-us-east-1/{name}_{version}_{build}.ht\"\n",
    "                    },\n",
    "                    \"gcp\": {\n",
    "                        \"us\": f\"gs://hail-datasets-us/{name}_{version}_{build}.ht\",\n",
    "                        \"eu\": f\"gs://hail-datasets-eu/{name}_{version}_{build}.ht\"\n",
    "                    }\n",
    "                },\n",
    "                \"version\": version\n",
    "            }\n",
    "        ]\n",
    "    }\n",
    "    datasets[name] = json_entry\n",
    "\n",
    "# Write new entries back to datasets.json config:\n",
    "with open(datasets_path, \"w\") as f:\n",
    "    json.dump(datasets, f, sort_keys=True, ensure_ascii=False, indent=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### WHR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Iterate through GIANT waist-hip ratio (BMI adj.) text files to write Hail Tables to GCS and generate entries to insert into datasets.json:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For GIANT 2018 Exome Array Summary Statistics - WHR\n",
    "sexes = [\"C\", \"M\", \"W\"]\n",
    "populations = [\"All\", \"Eur\"]\n",
    "renamed_populations = [\"ALL\", \"EUR\"]\n",
    "models = [\"Rec\", \"Add\"]\n",
    "\n",
    "# Generate valid combinations for existing filenames, e.g. ('C', 'All', 'Rec')\n",
    "combinations = [(sexes[i], populations[j], models[k]) \n",
    "                for i in range(0, len(sexes))\n",
    "                for j in range(0, len(populations))\n",
    "                for k in range(0, len(models))]\n",
    "\n",
    "# Map population name to maf and exac_maf field names\n",
    "maf = dict(zip(populations, [\"gmaf\", \"eur_maf\"]))\n",
    "exac_maf = dict(zip(populations, [\"exac_maf\", \"exac_nfe_maf\"]))\n",
    "\n",
    "for triplet in combinations:\n",
    "    print(triplet)\n",
    "    \n",
    "    input_file = os.path.expanduser(\"~\") + f\"/tmp/giant_whr_exome_summary/PublicRelease.WHRadjBMI.{triplet[0]}.{triplet[1]}.{triplet[2]}.txt\"\n",
    "    name = f\"giant_whr_exome_{triplet[0]}_{triplet[1].upper()}_{triplet[2]}\"\n",
    "    version = \"2018\"\n",
    "    build = \"GRCh37\"\n",
    "    \n",
    "    ht = hl.import_table(input_file,\n",
    "                         impute=True,\n",
    "                         missing=[\"-\", \"NA\", \"Inf\", \"\"],\n",
    "                         types = {\"beta\": hl.tfloat64,\n",
    "                                  \"se\": hl.tfloat64,\n",
    "                                  \"pvalue\": hl.tfloat64})\n",
    "\n",
    "    ht2 = ht.annotate(locus = hl.locus(ht.markername.split(\":\")[0],\n",
    "                                       hl.int(ht.markername.split(\":\")[1]),\n",
    "                                       reference_genome=build),\n",
    "                      alleles = [ht.ref, ht.alt],\n",
    "                      temp_maf = hl.dict(\n",
    "                          ht[maf[triplet[1]]].split(\",\").map(\n",
    "                              lambda x: (x.split(\":\")[0], hl.float(x.split(\":\")[1]))\n",
    "                          )\n",
    "                      ),\n",
    "                      temp_exac_maf = hl.dict(\n",
    "                          ht[exac_maf[triplet[1]]].split(\",\").map(\n",
    "                              lambda x: (x.split(\":\")[0], hl.float(x.split(\":\")[1]))\n",
    "                          )\n",
    "                      )\n",
    "                     )\n",
    "                  \n",
    "    ht2 = ht2.select(\"locus\", \"alleles\", \"snpname\", \"temp_maf\", \"temp_exac_maf\", \"beta\", \"se\", \"pvalue\", \"n\")\n",
    "    ht2 = ht2.rename({\"n\": \"sample_size\",\n",
    "                      \"snpname\": \"snp_name\",\n",
    "                      \"temp_maf\": maf[triplet[1]],\n",
    "                      \"temp_exac_maf\": exac_maf[triplet[1]]})\n",
    "    ht2 = ht2.key_by(\"locus\", \"alleles\")\n",
    "\n",
    "    n_rows = ht2.count()\n",
    "    n_partitions = ht2.n_partitions()\n",
    "\n",
    "    ht2 = ht2.annotate_globals(metadata=hl.struct(name=name,\n",
    "                                                  version=version,\n",
    "                                                  reference_genome=build,\n",
    "                                                  n_rows=n_rows,\n",
    "                                                  n_partitions=n_partitions))\n",
    "\n",
    "    for region in [\"us\", \"eu\"]:\n",
    "        output_file = f\"gs://hail-datasets-{region}/{name}_{version}_{build}.ht\"\n",
    "        ht2.write(output_file, overwrite=True)\n",
    "\n",
    "    json_entry = {\n",
    "        \"annotation_db\": {\n",
    "            \"key_properties\": [\n",
    "                \"unique\"\n",
    "            ]\n",
    "        },\n",
    "        \"description\": f\"GIANT (Genetic Investigation of ANthropometric Traits): \"\n",
    "                       f\"waist-hip ratio (WHR) adjusted for BMI exome array summary statistics Hail Table. \"\n",
    "                       f\"Note the following abbreviations used in filenames: C-combined sexes, M-men, W-women, \"\n",
    "                       f\"ALL-all ancestries, EUR-european descent only, Add-additive genetic model, \"\n",
    "                       f\"and Rec-recessive genetic model.\",\n",
    "        \"url\": \"https://portals.broadinstitute.org/collaboration/giant/index.php/GIANT_consortium_data_files\",\n",
    "        \"versions\": [\n",
    "            {\n",
    "                \"reference_genome\": build,\n",
    "                \"url\": {\n",
    "                    \"aws\": {\n",
    "                        \"us\": f\"s3://hail-datasets-us-east-1/{name}_{version}_{build}.ht\"\n",
    "                    },\n",
    "                    \"gcp\": {\n",
    "                        \"us\": f\"gs://hail-datasets-us/{name}_{version}_{build}.ht\",\n",
    "                        \"eu\": f\"gs://hail-datasets-eu/{name}_{version}_{build}.ht\"\n",
    "                    }\n",
    "                },\n",
    "                \"version\": version\n",
    "            }\n",
    "        ]\n",
    "    }\n",
    "    datasets[name] = json_entry\n",
    "\n",
    "# Write new entries back to datasets.json config:\n",
    "with open(datasets_path, \"w\") as f:\n",
    "    json.dump(datasets, f, sort_keys=True, ensure_ascii=False, indent=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Docs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate .rst files for dataset schemas:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import textwrap\n",
    "\n",
    "output_dir = os.path.abspath(\"../../hail/python/hail/docs/datasets/schemas\")\n",
    "datasets_path = os.path.abspath(\"../../hail/python/hail/experimental/datasets.json\")\n",
    "with open(datasets_path, \"r\") as f:\n",
    "    datasets = json.load(f)\n",
    "\n",
    "names = [name for name in list(datasets.keys()) if \"giant\" in name]\n",
    "for name in names:\n",
    "    versions = sorted(set(dataset[\"version\"] for dataset in datasets[name][\"versions\"]))\n",
    "    if not versions:\n",
    "        versions = [None]\n",
    "    reference_genomes = sorted(set(dataset[\"reference_genome\"] for dataset in datasets[name][\"versions\"]))\n",
    "    if not reference_genomes:\n",
    "        reference_genomes = [None]\n",
    "\n",
    "    print(name)\n",
    "    print(versions[0])\n",
    "    print(reference_genomes[0])\n",
    "\n",
    "    path = [dataset[\"url\"][\"gcp\"][\"us\"]\n",
    "            for dataset in datasets[name][\"versions\"]\n",
    "            if all([dataset[\"version\"] == versions[0],\n",
    "                    dataset[\"reference_genome\"] == reference_genomes[0]])]\n",
    "    assert len(path) == 1\n",
    "    path = path[0]\n",
    "\n",
    "    table = hl.methods.read_table(path)\n",
    "    description = table.describe(handler=lambda x: str(x))\n",
    "\n",
    "    if path.endswith(\".ht\"):\n",
    "        table_class = \"hail.Table\"\n",
    "    else:\n",
    "        table_class = \"hail.MatrixTable\"\n",
    "\n",
    "    template = \"\"\".. _{dataset}:\n",
    "\n",
    "{dataset}\n",
    "{underline1}\n",
    "\n",
    "*  **Versions:** {versions}\n",
    "*  **Reference genome builds:** {ref_genomes}\n",
    "*  **Type:** :class:`{class}`\n",
    "\n",
    "Schema ({version0}, {ref_genome0})\n",
    "{underline2}\n",
    "\n",
    ".. code-block:: text\n",
    "\n",
    "{schema}\n",
    "\n",
    "\"\"\"\n",
    "    context = {\n",
    "        \"dataset\": name,\n",
    "        \"underline1\": len(name) * \"=\",\n",
    "        \"version0\": versions[0],\n",
    "        \"ref_genome0\": reference_genomes[0],\n",
    "        \"versions\": \", \".join([str(version) for version in versions]),\n",
    "        \"ref_genomes\": \", \".join([str(reference_genome) for reference_genome in reference_genomes]),\n",
    "        \"underline2\": len(\"\".join([\"Schema (\", str(versions[0]), \", \", str(reference_genomes[0]), \")\"])) * \"~\",\n",
    "        \"schema\": textwrap.indent(description, \"    \"),\n",
    "        \"class\": table_class\n",
    "     }\n",
    "    with open(output_dir + \"/\" + name + \".rst\", \"w\") as f:\n",
    "        f.write(template.format(**context))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
